import yaml
import numpy as np
import mpmath as mp
from scipy.spatial.transform import Rotation

from .simple_sate_form import simple_sate_form
from ..file_io import global_output

class sate_motion(object):
    def __init__(self, microthrust, sate_mass) -> None:
        '''
        计算卫星运动所需时间及速度增量
        microthrust: 微推推力，单位\muN
        sate_mass: 卫星质量，单位kg
        '''
        self.__microthrust = microthrust
        self.__sate_mass = sate_mass

        self.__max_acceleration = self.__microthrust / self.__sate_mass

    def linear_motion(self, start_point, end_point, axis=0):
        '''
        L = 0.5 * a * t^2
        "Bang-bang"机动模式
        '''
        distance = np.linalg.norm(end_point-start_point, axis=axis)
        #delta_t = np.sqrt(distance / self.__max_acceleration)
        #delta_v = delta_t * self.__max_acceleration
        delta_v = np.sqrt(self.__microthrust * distance / (9. * self.__sate_mass))
        delta_t = 10 * distance / (9*delta_v)

        return delta_t / 3600. / 24., delta_v

    def circling_motion(self, start_point, period=np.pi*2., repetition=1):
        rotate = Rotation.from_rotvec(period * np.array([0., 0., 1.], dtype=np.float64))
        end_point = rotate.apply(start_point)
        delta_t, delta_v = self.linear_motion(start_point, end_point)

        if repetition != 1:
            rotate = Rotation.from_rotvec(period * repetition * np.array([0., 0., 1.], dtype=np.float64))
            end_point = rotate.apply(start_point)
            end_phase = period * repetition

        return delta_t*repetition, delta_v*repetition, end_point, end_phase

class sate_motion_circle(object):
    def __init__(self, microthrust, sate_mass) -> None:
        '''
        计算卫星运动所需时间及速度增量
        microthrust: 微推推力，单位\muN
        sate_mass: 卫星质量，单位kg
        '''
        self.__microthrust = microthrust
        self.__sate_mass = sate_mass

        self.__max_acceleration = self.__microthrust / self.__sate_mass
        self.__max_jac = self.get_max_JacobiSN()

        # 计算匀速圆周运动阶段总角度
        from scipy.integrate import quad
        radius = 30.
        jac_parm = np.sqrt(self.__microthrust / (self.__sate_mass * radius))
        theta_acc, _ = quad(self.get_omega, 0., self.__max_jac/jac_parm, args=(jac_parm))
        self.__theta_unf = 2*np.pi - 2*theta_acc

    def linear_motion(self, start_point, end_point, axis=0):
        '''
        L = 0.5 * a * t^2
        "Bang-bang"机动模式
        '''
        distance = np.linalg.norm(end_point-start_point, axis=axis)
        delta_t = np.sqrt(distance / self.__max_acceleration)
        delta_v = delta_t * self.__max_acceleration

        return delta_t / 3600. / 24., delta_v

    @staticmethod
    def get_max_JacobiSN():
        sn = mp.ellipfun('sn')
        def sn_scalar(u,m):
            return -np.abs(sn(np.linalg.norm(u), m).real)

        from scipy.optimize import minimize
        #sn_result = minimize(fun=sn_scalar, x0=(1.31), args=(-1), method = "Nelder-Mead", bounds=((0., 2.4),))
        sn_result = minimize(fun=sn_scalar, x0=(1.31), args=(-1), method = "L-BFGS-B", bounds=((0., 2.4),))
        return sn_result.x[0]

    @staticmethod
    def get_omega(time, jac_parm):
        sn = mp.ellipfun('sn')
        return jac_parm * sn(jac_parm*time, -1).real

    def circling_motion(self, start_point, period=np.pi*2., repetition=1):
        radius = np.linalg.norm(start_point)
        jac_parm = np.sqrt(self.__microthrust / (self.__sate_mass * radius))

        delta_t = 2*(self.__max_jac/jac_parm) + self.__theta_unf / jac_parm
        delta_v = self.__max_acceleration * delta_t

        delta_t = delta_t / 3600. / 24.

        return delta_t, delta_v, start_point, 0.




class relative_sate_form(simple_sate_form):
    '''
    基于Darwin运动学模型的卫星编队建模
    参考文献《Micropropulsion Technology Assessment for DARWIN》
    仅考虑四个子望远镜相对于合束器的相对运动。
    合束器视作静止物体，无速度增量。
    '''
    def __init__(self, yaml_config):
        #simple_sate_form.__init__(self, yaml_config)
        super().__init__(yaml_config)
        self.__microthrust  = yaml_config.get("MicroThrust", 10.0) * 1e-3 # 望远镜推进器推力 / mN
        self.__sate_mass    = yaml_config.get("SatelliteMass", 5e3)      # 望远镜卫星质量 / kg
        self.__form_phase   = yaml_config.get("InitialPhase", 0.)        # 初始编队相位 / rad
        self.__imaging_nulling_ratio = yaml_config.get("ImagingNullingRatio", 6.) # 望远镜成像基线与消零基线的比例
        self.__form_theta = np.arctan(1./self.__imaging_nulling_ratio)
        self.__form_half_diagonal = yaml_config.get("InitialBaseline", 10.) / np.sin(self.__form_theta) / 2.   # 初始消零基线 / m
        self.__delta_v_log_file = yaml_config.get("SaveDeltaV", True) # 是否存储DeltaV变化log文件

        self.__obs_repetition = yaml_config.get("ObsRepetition", 6) # 每圈观测取点数目 / 个
        self.__obs_rotation_radian = 2*np.pi / float(self.__obs_repetition)
        self.__obs_delta_mas = yaml_config.get("ObsDeltaMAS", 20.)  / 1e3 / 3600. / 180. * np.pi  # 相邻基线间隔毫弧度数 / mas

        self.__motion = sate_motion(self.__microthrust, self.__sate_mass)
        #self.__motion = sate_motion_circle(self.__microthrust, self.__sate_mass)
        self.__rotate_index = np.array([0, 1, 1, 0], dtype=np.int8)
        self.__theta_index = np.array([1, -1, 1, -1], dtype=np.int8)

        self.__sate_delta_v = np.zeros(4, dtype=np.float64)

        # 创建记录各星delta V变化的文件
        self.__delta_v_csv = None
        if self.__delta_v_log_file == True:
            self.__delta_v_csv = global_output.create_csv_file('delta_v_log')
            self.__delta_v_csv.write_head(["ID", "Sate1", "Sate2", "Sate3", "Sate4"])

    def update_satellite(self, current_target=[], total_time_cost = 0., sleep = False):
        if sleep == True:
            #simple_sate_form.orbit_lon = super().orbit_lon + super().sleep_time * super().orbit_rotate_angular_velocity
            #super().orbit_lon = super().orbit_lon + super().sleep_time * super().orbit_rotate_angular_velocity
            super().update_orbit_lon(super().sleep_time * super().orbit_rotate_angular_velocity)
            return -1, super().sleep_time

        # 计算各星目标变换状态
        self.__form_phase, _, delta_v = self.trace_planning(coord_star=[current_target[1], current_target[2]], \
                                                half_diagonal=self.__half_dia_determine(current_target[8], current_target[9]))
        self.__sate_delta_v += delta_v + current_target[11]
        if self.__delta_v_csv != None:
            self.__delta_v_csv.write_line([current_target[0]]+list(self.__sate_delta_v))

        # 更新卫星编队状态
        simple_sate_form.point_lon = current_target[1]
        simple_sate_form.point_lat = current_target[2]
        #simple_sate_form.orbit_lon = super().orbit_lon + total_time_cost * super().orbit_rotate_angular_velocity
        super().update_orbit_lon(total_time_cost * super().orbit_rotate_angular_velocity)
        self.__update_form_phase(current_target[12])

        # 更新最终基线状态
        #self.__form_half_diagonal = current_target[8]
        half_diagonal_mid = (current_target[8] + current_target[9]) / 2.
        if self.__form_half_diagonal < half_diagonal_mid:
            self.__form_half_diagonal = current_target[9]
        else:
            self.__form_half_diagonal = current_target[8]


        return current_target[0], total_time_cost

    def end_operation(self):
        global_output.write_log(f"Satellite delta v:")
        for i, deltav in enumerate(self.__sate_delta_v):
            global_output.write_log(f"Satellite {i}: {deltav} m/s")
        global_output.write_log(f"Satellites total delta v: {np.sum(self.__sate_delta_v)} m/s")
        self.__delta_v_csv.save()


    def observe_model(self, hz_inner, hz_outer, t_obs, mir_wave_length):
        form_phase = 0.
        hz_now = hz_outer

        #half_diagonal = 1.22 * mir_wave_length / (hz_now * np.sin(self.__form_theta) * 2.)
        half_diagonal = 0.59 * mir_wave_length / (hz_now * np.sin(self.__form_theta) * 2.)
        start_point = self.sate_in_origin_coord(half_diagonal, form_phase)[0]
        total_delta_t, total_delta_v, end_point, end_phase = \
            self.__motion.circling_motion(start_point, period=self.__obs_rotation_radian, repetition=self.__obs_repetition - 1)

        while hz_now > hz_inner:
            hz_now -= self.__obs_delta_mas
            #half_diagonal = 1.22 * mir_wave_length / (hz_now * np.sin(self.__form_theta) * 2.)
            half_diagonal = 0.59 * mir_wave_length / (hz_now * np.sin(self.__form_theta) * 2.)
            #print(f"{half_diagonal}")

            form_phase = form_phase + end_phase
            if form_phase > 2*np.pi:
                form_phase -= 2*np.pi
            start_point = self.sate_in_origin_coord(half_diagonal, form_phase)[0]

            line_delta_t, line_delta_v = self.__motion.linear_motion(start_point=end_point, end_point=start_point)
            cir_delta_t, cir_delta_v, end_point, end_phase = \
                self.__motion.circling_motion(start_point, period=self.__obs_rotation_radian, repetition=self.__obs_repetition - 1)

            total_delta_t = total_delta_t + line_delta_t + cir_delta_t + t_obs
            total_delta_v = total_delta_v + line_delta_v + cir_delta_v
        
        return total_delta_t, total_delta_v, form_phase, half_diagonal

    @staticmethod
    def coord_rotation_matrix(ra, dec):
        '''
        以J2000坐标系为基准，计算赤经赤纬为(ra, dec)恒星坐标系的旋转矩阵
        定义z轴为以合束器为原点，指向观测恒星的单位向量z_unit
        x轴x_unit由z_unit叉乘(0., 0., 1.)得到
        y轴y_unit由z_unit叉乘x_unit得到
        单位向量按列向量排列，形成的3x3矩阵即为旋转矩阵。
        '''
        if dec == 90.:
            print("Are you sure you want to detect the star with dec=90?")
            return np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        ra = ra / 180. * np.pi
        dec = dec / 180. * np.pi
        ref_z_unit = np.array([0., 0., 1.], dtype=np.float64)
        z_unit = np.array([np.cos(dec) * np.cos(ra), np.cos(dec) * np.sin(ra), np.sin(dec)], dtype=np.float64)
        x_unit = np.cross(z_unit, ref_z_unit)
        y_unit = np.cross(z_unit, x_unit)

        x_unit = x_unit / np.linalg.norm(x_unit)
        y_unit = y_unit / np.linalg.norm(y_unit)

        return np.mat(np.array([x_unit, y_unit, z_unit], dtype=np.float64)).T
        #return np.transpose(np.array([x_unit, y_unit, z_unit], dtype=np.float64))

    def sate_in_origin_coord(self, half_diagonal = 30., form_phase = 0.):
        #if type(form_phase) != float:
        #    form_phase = np.linalg.norm(form_phase)
        points_in_coord = np.zeros([4, 3], dtype=np.float64)
        radian = self.__theta_index*self.__form_theta + self.__rotate_index*np.pi + form_phase
        for i in range(4):
            points_in_coord[i,0] = half_diagonal * np.cos(radian[i])
            points_in_coord[i,1] = half_diagonal * np.sin(radian[i])

        return points_in_coord

    def coord_transform(self, coord_star=[350., -30]):

        #rotate_matrix1 = self.coord_rotation_matrix(self.__point_lon, self.__point_lat)
        rotate_matrix1 = self.coord_rotation_matrix(super().point_lon, super().point_lat)
        rotate_matrix2 = self.coord_rotation_matrix(coord_star[0], coord_star[1])

        origin_sate_points = self.sate_in_origin_coord(self.__form_half_diagonal, self.__form_phase)

        rotate1 = Rotation.from_matrix(rotate_matrix1)
        rotate2 = Rotation.from_matrix(rotate_matrix2)
        rotate = rotate2.inv() * rotate1

        sate_after_rotate = rotate.apply(origin_sate_points)
        return sate_after_rotate

    def sum_distance(self, form_phase:np.float64, sate_coords:np.ndarray, half_diagonal:np.float64):
        new_sate_coords = self.sate_in_origin_coord(half_diagonal=half_diagonal, form_phase=form_phase)
        distance_vec = np.linalg.norm(new_sate_coords-sate_coords, axis=1)
        return np.sum(distance_vec)

    def trace_planning(self, coord_star=[350., -30], half_diagonal=50):

        sate_coord = self.coord_transform(coord_star=coord_star)
        self.sum_distance(np.pi/8., sate_coord, half_diagonal)

        from scipy.optimize import minimize_scalar
        form_result = minimize_scalar(fun=self.sum_distance, args=(sate_coord, half_diagonal), \
                               method = "Bounded", bounds=(0., 2*np.pi))
        new_sate_coords = self.sate_in_origin_coord(half_diagonal=half_diagonal, form_phase=form_result.x)
        delta_t, delta_v = self.__motion.linear_motion(sate_coord, new_sate_coords, axis=1)
        return form_result.x, delta_t, delta_v

    def transform_cost(self, coord_star=[0., 0.], half_diagonal_inner=10., half_diagonal_outer=30.):
        phase, delta_t, delta_v = self.trace_planning(coord_star, self.__half_dia_determine(half_diagonal_inner, half_diagonal_outer))

        return phase, np.max(delta_t), np.sum(delta_v)

    @property
    def motion(self):
        return self.__motion

    @property
    def form_theta(self):
        return self.__form_theta

    def __half_dia_determine(self, half_dia_inn, half_dia_out):
        half_diagonal_mid = (half_dia_inn + half_dia_out) / 2.
        if self.__form_half_diagonal < half_diagonal_mid:
            return half_dia_inn
        else:
            return half_dia_out

    def __update_form_phase(self, phase):
        self.__form_phase += phase
        if self.__form_phase > 2*np.pi:
            self.__form_phase -= 2*np.pi
        elif self.__form_phase < 0:
            self.__form_phase += 2*np.pi